from fastapi import FastAPI, HTTPException, Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import requests
import json
from typing import Optional, Dict, Any, List

app = FastAPI()

# 配置 CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

OLLAMA_BASE_URL = "http://localhost:11434"

class GenerateRequest(BaseModel):
    model: str
    prompt: str
    stream: bool = False
    options: Optional[Dict[str, Any]] = None

class ChatRequest(BaseModel):
    model: str
    messages: List[Dict[str, str]]
    stream: bool = False
    format: Optional[str] = None
    options: Optional[Dict[str, Any]] = None
    template: Optional[str] = None

@app.post("/api/chat")
async def chat(request: ChatRequest):
    try:
        # Print request data for debugging
        print("Request data:", request.dict())
        
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/chat",
            json=request.dict(exclude_unset=True)  # Changed from exclude_none to exclude_unset
        )
        response.raise_for_status()

        if request.stream:
            return Response(
                content=response.content,
                media_type="text/event-stream",
                headers=dict(response.headers)
            )
        print("Response Header: " + str(response.headers))
        return Response(
            content=response.content,
            media_type="application/json",
            headers=dict(response.headers)
        )

    except requests.exceptions.RequestException as e:
        error_detail = str(e)
        if hasattr(e.response, 'text'):
            error_detail = e.response.text
            print("Error response:", error_detail)  # Print error response for debugging
        raise HTTPException(
            status_code=e.response.status_code if hasattr(e, 'response') else 500,
            detail=error_detail
        )

@app.post("/api/generate")
async def generate(request: GenerateRequest):
    try:
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/generate",
            json=request.dict(exclude_none=True)
        )
        response.raise_for_status()

        if request.stream:
            return Response(
                content=response.content,
                media_type="text/event-stream"
            )
        
        return Response(
            content=response.content,
            media_type="application/json"
        )

    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/version")
async def get_version():
    try:
        response = requests.get(f"{OLLAMA_BASE_URL}/api/version")
        response.raise_for_status()
        
        return Response(
            content=response.content,
            media_type="application/json; charset=utf-8",
            headers=dict(response.headers)
        )
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/tags")
async def list_models():
    try:
        response = requests.get(f"{OLLAMA_BASE_URL}/api/tags")
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/tags/{model_name}")
async def get_model_info(model_name: str):
    try:
        response = requests.get(f"{OLLAMA_BASE_URL}/api/tags/{model_name}")
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)